#coding=utf8

from numpy import *
from numpy.random import *
from matplotlib.pyplot import *

def sphere(p):
    temp = p ** 2
    return transpose([temp.sum(axis = 1)])

def rastrigin(p):
    t1 = p ** 2
    t2 = 10 * cos(2 * pi * p)
    t3 = t1 - t2 + 10
    t3 = t3.sum(axis = 1)
    return transpose([t3])

def func(x, y):
    return y + 1

def griewank(p):
    ind = fromfunction(func, (pop_size, dim), dtype = 'int')
    t1 = p ** 2
    t1 = t1.sum() / 4000.0
    t2 = cos(p / sqrt(ind))
    t2 = prod(t2, axis = 1)
    return transpose([t1 - t2 + 1])

#通过给foo赋值为不同的函数，来进行不同的分析
foo = None

#迭代次数
iter_max = 20000

#种群规模
pop_size = 30

#维度
dim = 20

#最大位置绝对值，即粒子的位置必须处于-100到100之间
pmax = 100

#最大速度绝对值，即粒子的速度必须处于-0.1pmax到0.1pmax之间
vmax = pmax * 0.1

#实验次数，多次实验取平均值可使得结果更加准确
test_times = 1

#d_low
d_low = 5.0e-6
d_high = 0.25

mode = 1

#初始位移
init_p = (rand(pop_size, dim) * 2.0 - 1.0) * pmax

#初始速度
init_v = (rand(pop_size, dim) * 2.0 - 1.0) * vmax

#这是一个基本PSO的实现，w取0.5，c1和c2均取2
class AbstractPSO():
    def __init__(self):
        self.refresh()
    
    def refresh(self):
        #初始化种群密度，每次迭代的密度值都被保存在里面
        self.diversity = zeros((iter_max, 1))

        #每次迭代的全局最优结果数组
        self.gbest_value = zeros((iter_max, 1))
        
        #每个粒子在每一维上的个人最优位移
        self.pbest = zeros((pop_size, dim))

        #全局最优位移
        self.gbest = zeros((pop_size, dim))
        
        #当前迭代的个人最优结果数组
        self.pbest_value = zeros((pop_size, 1))
        
    #求种群密度函数
    def diverse(self, p):
        L = sqrt(dim * ((2 * pmax) ** 2))
        S = average(p, axis = 0)
        t2 = 0.0
        for i in range(pop_size):
            t1 = 0.0
            for j in range(dim):
                t1 = t1 + (p[i, j] - S[j]) ** 2
            t2 = t2 + sqrt(t1)
        return t2 / L / pop_size  
   
    #粒子群算法启动   
    def run(self):
        
        self.refresh()
        
        #将全局的初始位移和初始速度复制一份到局部，以便于处理
        p = copy(init_p)
        v = copy(init_v)
        
        #获取当前目标函数处理后的值
        value = foo(p)
        
        #最优值为value数组中的最小值
        best_value = value.min()

        #最优值在数组中的位置
        best_arg = value.argmin()
        
        #将初始最优值保存到全局最优值数组中去
        self.gbest_value[0,0] = best_value
        
        #当前个人最优位移为初始位移
        self.pbest = p
        
        #当前个人最优值为value数组
        self.pbest_value = value
        
        #将所有粒子的全局最优位移都设置为当前的最佳位移
        for i in range(pop_size):
            self.gbest[i] = p[best_arg]
        
        #保存当前种群密度
        self.diversity[0, 0] = self.diverse(p)

        #位置更新
        p = p + v
        
        #迭代开始
        for i in range(1, iter_max):

            #算出当前种群密度值并保存
            self.diversity[i, 0] = self.diverse(p)
                        
            #得到目标函数的当前值
            value = foo(p)
            
            #获得最优值和最优值所在数组中的位置
            best_value = value.min()
            best_arg = value.argmin()
            
            #若当前value数组中的值小于个人最优值，则更新每个粒子的自身最优值和最优位移
            for j in range(pop_size):
                if value[j, 0] < self.pbest_value[j, 0]:
                    self.pbest_value[j, 0] = value[j, 0]
                    self.pbest[j] = p[j]

            #若value数组中的最优值比上次迭代的全局最优值还小，那么更新全局最优值，否则将上次迭代的全局最优值赋给这次 
            if best_value < self.gbest_value[i-1, 0]:
                self.gbest_value[i, 0] = best_value
                for j in range(pop_size):
                    self.gbest[j] = p[best_arg]
            else:
                self.gbest_value[i, 0] = self.gbest_value[i-1, 0]
            
            #获取w
            w = self.make_w(i)

            #获取c1和c2
            c1, c2 = self.make_c1_and_c2(i)

            #求两个随机数
            r1 = rand(pop_size, dim)
            r2 = rand(pop_size, dim)
            
            v = self.update(p, v, w, c1, c2, r1, r2, i)
            #限制速度
            for k in range(pop_size):
                for j in range(dim):
                    if v[k, j] > vmax:
                        v[k, j] = vmax
                    if v[k, j] < -vmax:
                        v[k, j] = -vmax
            
            #位移更新
            p = p + v
    
            #限制位移
            for k in range(pop_size):
                for j in range(dim):
                    if p[k, j] > pmax:
                        p[k, j] = pmax
                    if p[k, j] < -pmax:
                        p[k, j] = -pmax
        return self.gbest_value
    
    #c1和c2的生成器
    def make_c1_and_c2(self, iter):
        return 2.0, 2.0
    
    #w的生成器
    def make_w(self, iter):
        return 0.5
    
    def update(self, p, v, w, c1, c2, r1, r2, i):
        #更新速度值
        v = w * v + c1 * r1 * (self.pbest - p) + c2 * r2 * (self.gbest - p)

        return v
        
#w线性变换的PSO
class LWPSO(AbstractPSO):
    #将w进行线性变换 一开始w是0.9 通过递减之后变成0.4
    def make_w(self, iter):
        w_start = 0.9
        w_end = 0.4
        return w_start - (w_start - w_end) * double(iter) / iter_max 

#w随机变化的PSO
class RWPSO(AbstractPSO):
    def make_w(self, iter):
        return (0.8 + rand()) / 2.0 

#w反正切变化的PSO
class AWPSO(AbstractPSO):
    def make_w(self, iter):
        w_start = 0.9
        w_end = 0.4
        k = 0.4
        return (w_start - w_end) * arctan(1.56 * (1 - (double(iter) / iter_max) ** k)) + w_end

#c线性变化的PSO，c1和c2非对称
class LCPSO(AWPSO):
    def make_c1_and_c2(self, iter):
        c1_start = 2.75
        c1_end = 1.25
        c2_start = 0.5
        c2_end = 2.25
        
        c1 = c1_start + (c1_end - c1_start) * iter / iter_max
        c2 = c2_start + (c2_end - c2_start) * iter / iter_max
        
        return c1, c2

#c反余弦变化的PSO，c1和c2非对称
class ACPSO(AWPSO):
    def make_c1_and_c2(self, iter):
        c1_start = 2.75
        c1_end = 1.25
        c2_start = 0.5
        c2_end = 2.25
        
        c1 = c1_end + (c1_end - c1_start) * (arccos(-2.0 * double(iter) / iter_max + 1 ) / pi - 1.0)
        c2 = c2_end + (c2_end - c2_start) * (arccos(-2.0 * double(iter) / iter_max + 1 ) / pi - 1.0)

        return c1, c2

#生物启发式PSO，基于线性变化的c和w
class CEPSO(ACPSO):
    def update(self, p, v, w, c1, c2, r1, r2, i):
        global mode
                
        #更新速度值
        v = w * v + mode * ( c1 * r1 * (self.pbest - p) + c2 * r2 * (self.gbest - p))
        
        d1 = self.diversity[i-1, 0]
        d2 = self.diversity[i, 0]

        if (d1 < d_high and d1 >= d_low) and d2 < d_low :
            mode = -1.0
            v = -v
        elif (d1 < d_high and d1 >= d_low) and d2 > d_high:
            mode = 1.0
            v = -v

        return v
        
    
#获取用户对函数的选择,选择1就是Sphere函数，选择2就是Rastrigin函数，选择3就是Griewank函数
def fun_choose():
    print "Please choose a function for testing:"
    print "1. Sphere"
    print "2. Rastrigin"
    print "3. Griewank"
    return raw_input()

#w测试函数
def w_test():
    #声明foo是全局变量,不是局部变量
    global foo
    choice = fun_choose()
    #根据用户的选择，来选取不同的函数进行测试
    if choice == "1":
        foo = sphere
    elif choice == "2":
        foo = rastrigin
    elif choice == "3":
        foo = griewank
    else:
        print "illegal choice"
        return 

    #创建三种粒子群算法的对象
    pso1 = RWPSO()
    pso2 = AWPSO()
    pso3 = LWPSO()
    #value*用来保存三种粒子群算法的最优值变化
    value1 = 0.0
    value2 = 0.0
    value3 = 0.0
    #进行test_times次测试，把最优值累加
    for i in range(test_times):
        value1 = value1 + pso1.run()
        value2 = value2 + pso2.run()
        value3 = value3 + pso3.run() 
    #取最优值的加权平均之后，再取对数，将此图形保存到line*中
    line1 = plot(log10(value1 / (test_times)))
    line2 = plot(log10(value2 / (test_times)), ':')
    line3 = plot(log10(value3 / (test_times)), '--')
    #坐标图的标题
    title('w test')
    #x轴标签
    xlabel('times of iteration')
    #y轴标签
    ylabel('fitness(lg)')
    #在图上标注出图例
    legend((line1, line2, line3), ('random', 'arctan', 'linear'))
    #显示图象
    show()

#c测试函数
def c_test():
    #声明foo是全局变量,不是局部变量
    global foo
    choice = fun_choose()
    #根据用户的选择，来选取不同的函数进行测试
    if choice == "1":
        foo = sphere
    elif choice == "2":
        foo = rastrigin
    elif choice == "3":
        foo = griewank
    else:
        print "illegal choice"
        return 

    #创建二种粒子群算法的对象
    pso1 = LCPSO()
    pso2 = ACPSO()
    #value*用来保存二种粒子群算法的最优值变化
    value1 = 0.0
    value2 = 0.0
    #进行test_times次测试，把最优值累加
    for i in range(test_times):
        value1 = value1 + pso1.run()
        value2 = value2 + pso2.run()
    #取最优值的加权平均之后，再取对数，将此图形保存到line*中
    line1 = plot(log10(value1 / test_times))
    line2 = plot(log10(value2 / test_times), ':')
    #坐标图的标题
    title('c test')
    #x轴标签
    xlabel('times of iteration')
    #y轴标签
    ylabel('fitness(lg)')
    #在图上标注出图例
    legend((line1, line2), ('linear', 'arccos'))
    #显示图象
    show()

def ce_test():
    #声明foo是全局变量,不是局部变量
    global foo
    choice = fun_choose()
    #根据用户的选择，来选取不同的函数进行测试
    if choice == "1":
        foo = sphere
    elif choice == "2":
        foo = rastrigin
    elif choice == "3":
        foo = griewank
    else:
        print "illegal choice"
        return 

    pso1 = CEPSO()
    pso2 = ACPSO()
    
    value1 = pso1.gbest_value
    value2 = pso2.gbest_value
    
    d1 = pso1.diversity
    d2 = pso2.diversity
    
    for i in range(test_times):
        pso1.run()
        pso2.run()
        value1 = value1 + pso1.gbest_value
        value2 = value2 + pso2.gbest_value
        d1 = d1 + pso1.diversity
        d2 = d2 + pso2.diversity
        
    figure(1)
    subplot(211)
    line1 = plot(log10(value1 / test_times))
    line2 = plot(log10(value2 / test_times), ':')
    xlabel('times of iteration', fontsize = 10)
    ylabel('Fitness(lg)', fontsize = 10)
    legend((line1, line2), ('CE', 'AC'))
    
    subplot(212)
    line3 = plot(log10(d1 / test_times))
    line4 = plot(log10(d2 / test_times), ':')
    xlabel('times of iteration', fontsize = 10)
    ylabel('Diversity(lg)', fontsize = 10)
    legend((line3, line4), ('CE', 'AC'))
    show()
    
if __name__ == "__main__":
    print "Please choose a test case:"
    print "1.  w test"
    print "2.  c test"
    print "3.  ce test"
    choice = raw_input()
    if choice == "1":
        w_test()
    elif choice == "2":
        c_test()
    elif choice == "3":
        ce_test()
    else:
        print "illegal choice"
